/*
 *  Copyright CSCS 2012
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

#include <cusp/detail/device/generalized_spmv/csr_scalar.h>

#include <cusp/copy.h>
#include <cusp/array1d.h>
#include <cusp/exception.h>
#include <cusp/coo_matrix.h>
#include <cusp/csr_matrix.h>
#include <cusp/detail/random.h>
#include <cusp/detail/format_utils.h>

#include <cusp/print.h>

#include <thrust/copy.h>
#include <thrust/unique.h>
#include <thrust/set_operations.h>
#include <thrust/count.h>
#include <thrust/transform.h>
#include <thrust/transform_scan.h>
#include <thrust/iterator/constant_iterator.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/zip_iterator.h>

namespace cusp
{
namespace graph
{
namespace detail
{

// template <typename ValueType, typename MemorySpace>
//    template<typename MatrixType>
//    neighborhood_extension<ValueType,MemorySpace>
//    ::neighborhood_extension(const MatrixType& A, size_t nlevels)
//        : linear_operator<ValueType,MemorySpace>(A.num_rows, A.num_cols, A.num_rows)


template <typename Matrix>
Matrix extend_neighborhood(const Matrix& A_csr, size_t nlevels)
//                            cusp::csr_format, cusp::host_memory)
{
  typedef typename Matrix::index_type   IndexType;
  typedef typename Matrix::value_type   ValueType;
  typedef typename Matrix::memory_space MemorySpace;

  cusp::csr_matrix<IndexType,ValueType,cusp::host_memory> SP_csr(A_csr.num_rows, A_csr.num_cols, A_csr.num_entries );
  
  const IndexType N = A_csr.num_rows;
  size_t colcount = 0;
  size_t TotalSize = 0;

  int ColInds[N];       // temporary column indices
  int tmp[2*N];         // maximum possible size is 2*N
  int junk[N];          // intermediate 
  IndexType i;

  for ( i = 0; i < N; i++ )  // Over all rows
    {
      IndexType ioff_l = A_csr.row_offsets[i];
      IndexType ioff_h = A_csr.row_offsets[i+1];
      colcount =  ioff_h - ioff_l;
      thrust::copy(&A_csr.column_indices[ioff_l],&A_csr.column_indices[ioff_h],ColInds);
      thrust::copy(ColInds,&ColInds[colcount],junk);
      size_t current_size = colcount;
      for ( size_t lev=0; lev<nlevels; lev++ )
        {  // Expand ColInds by considering all neighbors
          for ( size_t j=0; j<colcount; j++) 
            {
              IndexType ind = ColInds[j];
              int joff_l = A_csr.row_offsets[ind]; 
              int joff_h = A_csr.row_offsets[ind+1]; 
              int *tmp_end = thrust::set_union(junk, &junk[current_size], 
                                               &A_csr.column_indices[joff_l],
                                               &A_csr.column_indices[joff_h], tmp );
              int *new_end = thrust::unique(tmp, tmp_end);
              thrust::copy(tmp,new_end,junk);
              current_size = new_end - tmp;
//              for(size_t i = 0; i < current_size; i++)
//                std::cout << "result[" << i << "] = " << junk[i] << std::endl;
            }
            colcount = current_size;
            thrust::copy(junk,&junk[current_size],ColInds);
        }
        // At this point ColInds should contain the multi-level neighbor set
        if ( TotalSize + current_size > SP_csr.num_entries ) 
          SP_csr.resize(A_csr.num_rows, A_csr.num_cols, 
                        SP_csr.num_entries+A_csr.num_entries );  // Add more non-zeros
        SP_csr.row_offsets[i] = TotalSize;  
        thrust::copy(ColInds,&ColInds[current_size],&SP_csr.column_indices[TotalSize]);
        TotalSize += current_size;
    }
    SP_csr.resize(A_csr.num_rows, A_csr.num_cols, TotalSize );  // shrink to fit
    SP_csr.row_offsets[N] = TotalSize;
    // cusp::print(SP_csr);
    return SP_csr;
}


//////////////////
// General Path //
//////////////////

template <typename Matrix, typename Format, typename MemorySpace>
Matrix neighborhood_extension(const Matrix& A, size_t level,
                              Format, MemorySpace)
{
  typedef typename Matrix::index_type   IndexType;
  typedef typename Matrix::value_type   ValueType;

  // convert matrix to CSR format and compute on the host
  cusp::csr_matrix<IndexType,ValueType,cusp::host_memory> A_csr(A);

  return cusp::graph::neighborhood_extension(A_csr, level);

}

} // end namespace detail

/////////////////
// Entry Point //
/////////////////

template <typename Matrix>
Matrix neighborhood_extension(const Matrix& A, size_t level)
{
    typedef typename Matrix::index_type   IndexType;
    typedef typename Matrix::value_type   ValueType;

    CUSP_PROFILE_SCOPED();

    if(A.num_rows != A.num_cols)
        throw cusp::invalid_input_exception("matrix must be square");

    if (level == 0)
    {
        return A;
    }
    else
    {
          cusp::csr_matrix<IndexType,ValueType,cusp::host_memory> SP_csr(A.num_rows, A.num_cols, A.num_entries );
 
//        return cusp::graph::detail::extend_neighborhood(A, level, typename Matrix::format(), typename Matrix::memory_space());
        SP_csr = cusp::graph::detail::extend_neighborhood(A, level);
        return SP_csr;
    }
}

} // end namespace graph
} // end namespace cusp

